import datetime
import pandas as pd
import tushare as ts
import matplotlib.pyplot as plt
from matplotlib.widgets import MultiCursor
from pandas.plotting import register_matplotlib_converters

import QUANTAXIS as QA
from QAStrategy.QAStrategy.classic_strategy.turtle_strategy import TurtleTrade


register_matplotlib_converters()


def main(codes, start, end):
    # print(codes.head())
    data = QA.QA_fetch_stock_day_adv(codes['code'].to_list(), start, end)
    data = data.to_qfq()
    ind = data.groupby(level=1, sort=False).apply(QA.QA_indicator_Turtle,
                                                              N_UP=55,
                                                              N_DN=20,
                                                              N_S_MA=10,
                                                              N_L_MA=60,
                                                              N_ATR=20)
    df_result = pd.DataFrame(columns=['code', 'time', 'CLOSE', 'UP', 'DIFF'])
    for _, code in codes.iterrows():
        # print(code)
        # df = data.data.loc[(slice(None), code), :]
        df = data.data.xs(code['code'], level='code', axis=0, drop_level=True)
        ind_last_bar = ind.iloc[-2]
        ind_current_bar = ind.iloc[-1]
        # if ind_current_bar.CLOSE > ind_last_bar.UP and ind_current_bar.DIFF > 0 and ind_current_bar.L_MA > 0:
        # if ind_current_bar.CLOSE > ind_last_bar.UP:
        if ind_current_bar.DIFF > 0:
            # and ind_current_bar.name[0].year == 2020 \
            # and ind_current_bar.name[0].month == 12 \
            # and ind_current_bar.name[0].day >= 1:
            print('code：{}---{}---时间{}---CLOSE{}---UP{}---DIFF{}'.format(
                code['code'], code['name'], ind_current_bar.name[0], ind_current_bar.CLOSE, ind_last_bar.UP, ind_current_bar.DIFF))
            df_result = df_result.append({'code': code['code'],
                                          'time': ind_current_bar.name[0],
                                          'CLOSE': ind_current_bar.CLOSE,
                                          'UP': ind_last_bar.UP,
                                          'DIFF': ind_current_bar.DIFF}, ignore_index=True)


            figure = plt.figure()
            plt.title('{}---'.format(code['code']))
            axes1 = figure.add_subplot(411)
            axes2 = figure.add_subplot(412)
            axes3 = figure.add_subplot(413)
            axes4 = figure.add_subplot(414)
            axes1.plot(df.index, df['CLOSE'])
            axes2.plot(df.index, df['CLOSE'])
            axes3.plot(df.index, df['CLOSE'])
            axes4.plot(df.index, df['CLOSE'])

            left, right = axes1.get_xlim()
            # axes1.hlines(y=box.high, xmin=left, xmax=right, linestyles='dashed')
            multi = MultiCursor(figure.canvas, (axes1, axes2, axes3, axes4), color='r', lw=1)

            plt.show()
    df_result.to_csv('turtle.csv')

    # data.to_csv('002581.csv')
    # print(data.index)


if __name__ == '__main__':
    ts_token = '17056d23a59ab71cb979c6a30185e092aba605c4544dac900a3eb7f8'
    ts.set_token(ts_token)
    pro = ts.pro_api()
    data = pro.stock_basic(exchange='', list_status='L', fields='symbol,name,area,industry,list_date')
    data = data[data.symbol.str.startswith('00') | data.symbol.str.startswith('60')]
    # data = data[data.symbol.str.startswith('002068')]
    data = data[data.list_date < '20160101']
    all = data.rename(columns={'symbol': 'code'})
    codes = all[~all.name.str.contains('ST')]
    # codes = codes.iloc[0: 1]
    # hs300 = pd.read_csv('D:\\PythonPro\\QUANTAXIS\\AI\\hs300.csv')
    # hs300 = ts.get_hs300s()
    # codes = codes['code'].to_list()
    # codes = ['002027']
    # print(codes)
    start = '2020-06-01'
    end = '2020-12-11'
    main(codes, start, end)
